#include <iostream>
#include<opencv2/opencv.hpp>
#include<math.h>
#include "yolo.h"

#include "parameter.h"

#include <memory>
#include <chrono>

// using namespace std;
using namespace cv;
using namespace dnn;

int main()
{	
	ParameterReader conf("../config.txt");
	std::string model_path = conf.getData("model_path").c_str();
	int height = atoi( conf.getData("height").c_str() );
	int width = atoi( conf.getData("width").c_str() );
	
	cv::VideoCapture cap(0);
	cap.set(cv::CAP_PROP_FRAME_WIDTH, width);
    cap.set(cv::CAP_PROP_FRAME_HEIGHT, height);

	// 定义Yolo，和Net类;
	std::shared_ptr<Yolo> test = std::make_shared<Yolo>();
	std::shared_ptr<Net> net = std::make_shared<Net>(); 
	
	// 读取网络
	if (test->readModel(*net, model_path, atoi( conf.getData("isCuda").c_str()) )) {
		std::cout << "--Load model ok!" << std::endl;
	}
	else {
		std::cout << "--Load model failed!" << std::endl;
		return -1;
	}
	std::cout << "\n-----------------" << std::endl;
	
	cv::Mat  frame, img_box;  // 图像输入,检查结果
	bool success_det;
	
	if(!cap.isOpened())
		{	
			std::cout << "--Please ckeck the VideoCapture" << std::endl;
			return -1;
		}
	while(cap.isOpened())
    {	
		std::vector<Output> result;   // 返回的结果,每次需要更新结果
		// cap >> img;
		// frame = img(cv::Rect(0, 0, 640, 480)).clone();
		cap >> frame;
		cv::resize(frame, frame, cv::Size(width, height));

		std::chrono::system_clock::time_point t1 = std::chrono::system_clock::now();
		success_det = test->Detect(frame, *net, result);
		auto t2 = std::chrono::system_clock::now();
		auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(t2 - t1);

		// std::cout <<  << std::endl;

		if(success_det)
		{
			img_box = test->drawPred(frame, result);
			// cv::resize(frame, frame, cv::Size(640, 480));
			cv::imshow("result", img_box);
		}
		else 
		{
			std::cout << "Detect Failed!"<< std::endl;
			cv::imshow("result", frame);
		}
		if(cv::waitKey(10) == 'q')
				break;
		
		// 打印结果
		if(atoi( conf.getData("cout_flag").c_str() ))
		{
			for(auto it:result)
			{
				std::cout  << test->class_select[it.id]<< "\t" 
							<< it.box.x << "\t" << it.box.y << "\t" 
							<< it.box.width << "\t" << it.box.height
							<< "\t" << it.confidence<< std::endl;
			}
		}
		std::cout<<"--Time cost = "<<1000/duration.count()<<" fps. "<<std::endl;
		std::cout << "----------" << std::endl;
		
    }
	
	return 0;
}

